#include "Neunet.hpp"

using namespace std;

// Position3D
Position3D::Position3D(double x, double y, double z) : x(x), y(y), z(z) {}
double Position3D::distance_to(const Position3D& other) const {
    double dx = x - other.x;
    double dy = y - other.y;
    double dz = z - other.z;
    return sqrt(dx*dx + dy*dy + dz*dz);
}

// Connection
Connection::Connection(SimpleNeuron* t, double w) : target(t), weight(w) {}

// 静态成员初始化
int SimpleNeuron::next_id = 1;
int SimpleNeuron::global_step = 0;

// ==================== 核心学习机制 ====================
/**
 * 接收输入刺激
 * @param strength 输入强度
 * @param source 输入来源神经元（可选，用于学习）
 */
void SimpleNeuron::receive_input(double strength, SimpleNeuron* source) {
    activation_level += strength;
    input_history.push_back(strength);
    if (input_history.size() > HISTORY_LENGTH) {
        input_history.pop_front();
    }
}

/**
 * 检查是否应该激活
 * 不同类型神经元有不同的激活条件
 */
bool SimpleNeuron::should_activate() const {
    // 检查不应期
    if (global_step - last_activation_step < REFRACTORY_STEPS) {
        return false;
    }

    // 根据神经元类型检查激活阈值
    switch (type) {
        case EXCITATORY:
            return activation_level >= ACTIVATION_THRESHOLD;
        case INHIBITORY:
            return activation_level >= INHIBITORY_THRESHOLD;
        case MODULATORY:
            return activation_level >= MODULATORY_THRESHOLD;
        default:
            return false;
    }
}

/**
 * 执行激活 - 不同类型神经元有不同的激活行为
 */
void SimpleNeuron::activate() {
    if (!should_activate()) return;

    // 记录激活
    activation_history.push_back(global_step);
    last_activation_step = global_step;

    // 重置激活水平
    activation_level = 0.0;

    // 根据神经元类型执行不同的激活行为
    switch (type) {
        case EXCITATORY:
            activate_excitatory();
            break;
        case INHIBITORY:
            activate_inhibitory();
            break;
        case MODULATORY:
            activate_modulatory();
            break;
    }

    // 限制激活历史长度
    if (activation_history.size() > HISTORY_LENGTH) {
        activation_history.pop_front();
    }
}

/**
 * 兴奋性神经元激活：正常的预测性信号传递
 */
void SimpleNeuron::activate_excitatory() {
    for (auto& conn : outputs) {
        if (conn->target) {
            conn->pre_activations.push_back(global_step);
            double signal = conn->weight * ACTIVATION_THRESHOLD;
            conn->target->receive_input(signal, this);
        }
    }
}

/**
 * 抑制性神经元激活：发送抑制信号，平衡网络
 */
void SimpleNeuron::activate_inhibitory() {
    for (auto& conn : outputs) {
        if (conn->target) {
            conn->pre_activations.push_back(global_step);
            // 抑制性信号：负向影响目标神经元
            double inhibitory_signal = -conn->weight * INHIBITORY_STRENGTH;
            conn->target->receive_input(inhibitory_signal, this);
        }
    }
}

/**
 * 调制性神经元激活：发送调制信号，影响整个区域
 */
void SimpleNeuron::activate_modulatory() {
    // 更新调制状态
    modulatory_state = 1.0;
    last_modulatory_event = global_step;

    // 调制性神经元影响范围更广，不仅限于直接连接
    for (auto& conn : outputs) {
        if (conn->target) {
            conn->pre_activations.push_back(global_step);

            // 根据当前情况发送不同类型的调制信号
            double modulation_signal = determine_modulation_signal();
            conn->target->receive_modulation(modulation_signal, this);
        }
    }
}

/**
 * 确定调制信号类型和强度
 */
double SimpleNeuron::determine_modulation_signal() {
    // 简化的调制逻辑：根据输入模式决定调制类型
    double recent_activity = 0.0;
    int recent_count = min(10, (int)input_history.size());
    for (int i = input_history.size() - recent_count; i < input_history.size(); ++i) {
        recent_activity += abs(input_history[i]);
    }

    if (recent_activity > 50.0) {
        return ALERT_SIGNAL;        // 高活动 -> 警报信号
    } else if (recent_activity > 20.0) {
        return ATTENTION_SIGNAL;    // 中等活动 -> 注意力信号
    } else {
        return REWARD_SIGNAL;       // 低活动 -> 奖励信号
    }
}

/**
 * 接收调制信号
 */
void SimpleNeuron::receive_modulation(double modulation, SimpleNeuron* source) {
    // 调制信号不直接影响激活水平，而是调整神经元的敏感性
    if (modulation > 0) {
        // 正向调制：提高敏感性
        activation_level += modulation * 0.5;

        // 调制信号还会影响学习速率
        if (type == EXCITATORY) {
            // 在接收到注意力信号时，增强学习
            for (auto& conn : outputs) {
                conn->weight *= 1.01;  // 微调权重
                conn->weight = min(MAX_WEIGHT, conn->weight);
            }
        }
    }
}

/**
 * 突触可塑性调整 - 只有兴奋性神经元进行STDP学习
 */
void SimpleNeuron::adjust_weights() {
    if (type != EXCITATORY) return;  // 只有兴奋性神经元进行STDP学习

    for (auto& conn : outputs) {
        if (!conn->target) continue;

        // 获取目标神经元的激活历史
        const auto& post_activations = conn->target->activation_history;

        // 对每个突触前激活，寻找相关的突触后激活
        for (int pre_step : conn->pre_activations) {
            for (int post_step : post_activations) {
                int time_diff = post_step - pre_step;

                // 在STDP窗口内才进行调整
                if (abs(time_diff) <= STDP_WINDOW) {
                    double weight_change = 0.0;

                    if (time_diff > 0 && time_diff <= STDP_WINDOW) {
                        // 突触前在前，突触后在后 -> 增强连接（因果关系）
                        weight_change = LEARNING_RATE * exp(-time_diff / 10.0);
                    } else if (time_diff < 0 && time_diff >= -STDP_WINDOW) {
                        // 突触后在前，突触前在后 -> 减弱连接（非因果）
                        weight_change = -LEARNING_RATE * exp(time_diff / 10.0);
                    }

                    // 应用权重变化
                    conn->weight += weight_change;
                    conn->weight = std::max(MIN_WEIGHT, std::min(MAX_WEIGHT, conn->weight));
                }
            }
        }

        // 清理过期的激活记录
        int cutoff_step = global_step - STDP_WINDOW;
        conn->pre_activations.erase(
            std::remove_if(conn->pre_activations.begin(), conn->pre_activations.end(),
                [cutoff_step](int step) { return step < cutoff_step; }),
            conn->pre_activations.end()
        );
    }
}

/**
 * 空间连接形成 - 不同类型神经元有不同的连接策略
 */
void SimpleNeuron::form_connections(std::vector<SimpleNeuron*>& all_neurons) {
    if (outputs.size() >= MAX_CONNECTIONS) return;

    double connection_range = (type == MODULATORY) ? MODULATORY_RANGE : CONNECTION_RADIUS;

    // 寻找连接范围内的神经元
    for (auto* target : all_neurons) {
        if (target == this) continue;
        if (outputs.size() >= MAX_CONNECTIONS) break;

        // 检查距离
        if (position.distance_to(target->position) > connection_range) continue;

        // 检查是否已连接
        bool already_connected = false;
        for (const auto& conn : outputs) {
            if (conn->target == target) {
                already_connected = true;
                break;
            }
        }

        // 不同类型神经元的连接偏好
        double connection_prob = get_connection_probability(target);

        // 以一定概率建立新连接
        if (!already_connected && dis(gen) < connection_prob) {
            double initial_weight = dis(gen) * 0.4 + 0.1; // 0.1-0.5之间随机权重
            outputs.push_back(std::make_unique<Connection>(target, initial_weight));
        }
    }
}

/**
 * 获取与目标神经元的连接概率
 */
double SimpleNeuron::get_connection_probability(SimpleNeuron* target) {
    switch (type) {
        case EXCITATORY:
            // 兴奋性神经元偏好连接到所有类型
            return CONNECTION_PROBABILITY;
        case INHIBITORY:
            // 抑制性神经元偏好连接到兴奋性神经元
            return (target->type == EXCITATORY) ? CONNECTION_PROBABILITY * 1.5 : CONNECTION_PROBABILITY * 0.5;
        case MODULATORY:
            // 调制性神经元倾向于广泛连接
            return CONNECTION_PROBABILITY * 0.8;
        default:
            return CONNECTION_PROBABILITY;
    }
}

/**
 * 移除弱连接 - 模拟神经修剪
 */
void SimpleNeuron::prune_weak_connections() {
    outputs.erase(
        std::remove_if(outputs.begin(), outputs.end(),
            [](const std::unique_ptr<Connection>& conn) {
                return conn->weight < PRUNE_THRESHOLD;
            }),
        outputs.end()
    );
}

/**
 * 每步更新 - 神经元的生命周期
 */
void SimpleNeuron::update() {
    // 激活水平自然衰减
    activation_level *= DECAY_RATE;

    // 调制状态衰减
    if (type == MODULATORY) {
        modulatory_state *= 0.95;
    }

    // 检查并执行激活
    if (should_activate()) {
        activate();
    }

    // 定期调整权重（仅兴奋性神经元）
    if (global_step % 10 == 0) {
        adjust_weights();
    }

    // 清理过期历史
    cleanup_history();
}

// ==================== 信息获取 ====================
struct NeuronInfo {
    int id;
    NeuronType type;
    double activation_level;
    int connections_count;
    double average_weight;
    int recent_activations;
    double modulatory_state;
};

SimpleNeuron::NeuronInfo SimpleNeuron::get_info() const {
    SimpleNeuron::NeuronInfo info;
    info.id = neuron_id;
    info.type = type;
    info.activation_level = activation_level;
    info.connections_count = outputs.size();
    info.modulatory_state = modulatory_state;

    // 计算平均权重
    if (!outputs.empty()) {
        double total_weight = 0.0;
        for (const auto& conn : outputs) {
            total_weight += conn->weight;
        }
        info.average_weight = total_weight / outputs.size();
    } else {
        info.average_weight = 0.0;
    }

    // 计算最近的激活次数
    int recent_cutoff = global_step - 50;
    info.recent_activations = std::count_if(
        activation_history.begin(), activation_history.end(),
        [recent_cutoff](int step) { return step >= recent_cutoff; }
    );

    return info;
}

// ==================== 静态方法 ====================
void SimpleNeuron::advance_global_step() { global_step++; }
int SimpleNeuron::get_global_step() { return global_step; }
void SimpleNeuron::reset_global_step() { global_step = 0; }

void SimpleNeuron::cleanup_history() {
    int cutoff_step = global_step - HISTORY_LENGTH;

    // 清理激活历史
    while (!activation_history.empty() && activation_history.front() < cutoff_step) {
        activation_history.pop_front();
    }

    // 清理输入历史（保持固定长度）
    while (input_history.size() > HISTORY_LENGTH) {
        input_history.pop_front();
    }
}

// 获取神经元类型名称
string SimpleNeuron::get_type_name() const {
    switch (type) {
        case EXCITATORY: return "兴奋性";
        case INHIBITORY: return "抑制性";
        case MODULATORY: return "调制性";
        default: return "未知";
    }
}

// ==================== 简化的神经网络类 ====================
SimpleNeuralNetwork::SimpleNeuralNetwork() = default;

/**
 * 添加神经元到网络
 */
void SimpleNeuralNetwork::add_neuron(double x, double y, double z, NeuronType type) {
    neurons.push_back(std::make_unique<SimpleNeuron>(x, y, z, type));
}

/**
 * 创建随机分布的神经网络，包含不同类型的神经元
 */
void SimpleNeuralNetwork::create_random_network(int num_neurons, double space_size) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis(0.0, space_size);
    std::uniform_real_distribution<> type_dis(0.0, 1.0);

    // 创建不同类型的神经元
    for (int i = 0; i < num_neurons; ++i) {
        double type_rand = type_dis(gen);
        NeuronType type;

        if (type_rand < 0.7) {
            type = EXCITATORY;      // 70% 兴奋性
        } else if (type_rand < 0.9) {
            type = INHIBITORY;      // 20% 抑制性
        } else {
            type = MODULATORY;      // 10% 调制性
        }

        add_neuron(dis(gen), dis(gen), dis(gen), type);
    }

    // 建立初始连接
    reorganize_connections();
}

/**
 * 重新组织连接
 */
void SimpleNeuralNetwork::reorganize_connections() {
    std::vector<SimpleNeuron*> neuron_ptrs;
    for (auto& neuron : neurons) {
        neuron_ptrs.push_back(neuron.get());
    }
    for (auto& neuron : neurons) {
        neuron->form_connections(neuron_ptrs);
        neuron->prune_weak_connections();
    }
}

/**
 * 仿真一步
 */
void SimpleNeuralNetwork::simulate_step() {
    SimpleNeuron::advance_global_step();
    for (auto& neuron : neurons) {
        neuron->update();
    }
    if (SimpleNeuron::get_global_step() % REORGANIZE_INTERVAL == 0) {
        reorganize_connections();
    }
}

/**
 * 给随机神经元输入刺激
 */
void SimpleNeuralNetwork::add_stimulus(double strength) {
    if (neurons.empty()) return;
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, neurons.size() - 1);
    int target_idx = dis(gen);
    neurons[target_idx]->receive_input(strength);
}

/**
 * 输出网络状态
 */
void SimpleNeuralNetwork::print_status() const {
    cout << "=== 神经网络状态 (步数: " << SimpleNeuron::get_global_step() << ") ===" << endl;

    int total_connections = 0;
    int total_activations = 0;
    double total_avg_weight = 0.0;
    int type_counts[3] = {0, 0, 0}; // 兴奋性、抑制性、调制性计数

    for (const auto& neuron : neurons) {
        auto info = neuron->get_info();
        total_connections += info.connections_count;
        total_activations += info.recent_activations;
        total_avg_weight += info.average_weight;
        type_counts[info.type]++;

        std::cout << neuron->get_type_name() << "神经元" << info.id 
                  << ": 激活=" << info.activation_level
                  << ", 连接=" << info.connections_count
                  << ", 权重=" << info.average_weight
                  << ", 近期激活=" << info.recent_activations;
                  
        if (info.type == MODULATORY) {
            cout << ", 调制状态=" << info.modulatory_state;
        }
        cout << endl;
    }
    
    cout << "网络统计: 平均连接=" << (double)total_connections/neurons.size()
              << ", 总激活=" << total_activations
              << ", 平均权重=" << total_avg_weight/neurons.size() << endl;
    cout << "神经元类型: 兴奋性=" << type_counts[0] 
              << ", 抑制性=" << type_counts[1] 
              << ", 调制性=" << type_counts[2] << endl;
    cout << endl;
}

size_t SimpleNeuralNetwork::size() const { return neurons.size(); }